import os
import numpy as np
import faiss
from util import createXY
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import argparse
import logging
from tqdm import tqdm

# 确保 FaissKNeighbors 类已经正确定义
from FaissKNeighbors import FaissKNeighbors

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def get_args():
    parser = argparse.ArgumentParser(description='使用CPU或GPU训练模型。')
    parser.add_argument('-m', '--mode', type=str, required=True, choices=['cpu', 'gpu'], help='选择训练模式：CPU或GPU。')
    parser.add_argument('-f', '--feature', type=str, required=True, choices=['flat', 'vgg'], help='选择特征提取方法：flat或vgg。')
    parser.add_argument('-l', '--library', type=str, required=True, choices=['sklearn', 'faiss'], help='选择使用的库：sklearn或faiss。')
    return parser.parse_args()

def main():
    args = get_args()
    logging.info(f"选择模式是 {args.mode.upper()}")
    logging.info(f"选择特征提取方法是 {args.feature.upper()}")
    logging.info(f"选择使用的库是 {args.library.upper()}")

    X, y = createXY(train_folder="test_set", dest_folder=".", method=args.feature)
    X = np.array(X).astype('float32')
    faiss.normalize_L2(X)
    y = np.array(y)
    logging.info("数据加载和预处理完成。")

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")

    best_k = -1
    best_accuracy = 0.0

    k_values = range(1, 6)

    KNNClass = FaissKNeighbors if args.library == 'faiss' else KNeighborsClassifier
    logging.info(f"使用的库为: {args.library.upper()}")

    for k in tqdm(k_values, desc='寻找最佳k值'):
        knn = KNNClass(k=k)  # 如果 FaissKNeighbors 需要其他参数，请添加
        knn.fit(X_train, y_train)
        accuracy = knn.score(X_test, y_test)
        if accuracy > best_accuracy:
            best_k = k
            best_accuracy = accuracy

    logging.info(f'最佳k值: {best_k}, 最高准确率: {best_accuracy}')

if __name__ == 'main__':
    main()